import sqlite3
import json
import signal
import threading
import random
import time
from datetime import datetime, timedelta
from wxauto import WeChat
import requests
from tenacity import retry, stop_after_attempt, wait_exponential

# 配置参数
TARGET_CHATS = ["测试群聊", "测试群2"]
DB_PATH = "wechat_group_messages.db"
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
GEMINI_API_KEY = "AIzaSyD91uzSkCR3bRUdpupNxhE4MbiDEeox0m0"  # ⚠️替换成你自己的 key

# 人性化模式设置
HUMAN_LIKE_MODE = True
MIN_DELAY = 1.0
MAX_DELAY = 5.0

# 批量处理设置
BATCH_PROCESS_INTERVAL = 3600  # 每1小时处理一次（秒）

# 全局变量
stop_event = threading.Event()
last_messages = {}


def human_like_delay():
    """添加随机延迟，模拟人类操作"""
    if HUMAN_LIKE_MODE:
        delay = random.uniform(MIN_DELAY, MAX_DELAY)
        time.sleep(delay)


def init_db(db_path: str):
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            chat_name TEXT NOT NULL,
            sender TEXT,
            content TEXT,
            msg_time TEXT NOT NULL,
            raw TEXT
        )
    """)
    cur.execute("""
        CREATE TABLE IF NOT EXISTS important_messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            chat_name TEXT NOT NULL,
            sender TEXT,
            content TEXT,
            msg_time TEXT NOT NULL,
            raw TEXT
        )
    """)
    cur.execute("CREATE INDEX IF NOT EXISTS idx_messages_chat_time ON messages(chat_name, msg_time)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_important_chat_time ON important_messages(chat_name, msg_time)")
    conn.commit()
    conn.close()


def safe_serialize(obj):
    """确保对象可被 JSON 序列化"""
    try:
        json.dumps(obj)
        return obj
    except Exception:
        return str(obj)


def save_message(chat_name: str, sender: str, content: str, msg_time: str, raw: dict, important=False):
    """保存消息到数据库"""
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()

    safe_raw = {k: safe_serialize(v) for k, v in raw.items()}
    safe_raw['important'] = important

    table = "important_messages" if important else "messages"

    try:
        cur.execute(f"""
            INSERT INTO {table} (chat_name, sender, content, msg_time, raw)
            VALUES (?, ?, ?, ?, ?)
        """, (chat_name, sender, content, msg_time, json.dumps(safe_raw, ensure_ascii=False)))
        conn.commit()
        tag = "重要" if important else "普通"
        print(f"[DB INFO] 成功保存{tag}消息：{chat_name}@{sender} -> {content}", flush=True)
    except Exception as e:
        print(f"[DB ERROR] 保存消息失败: {e}", flush=True)
    finally:
        conn.close()


# ============ AI 分析 ============
def analyze_message(content: str) -> bool:
    """
    调用 Gemini 模型，判断是否重要消息。
    返回 True 表示重要，False 表示普通。
    """
    try:
        headers = {
            "Content-Type": "application/json",
            "x-goog-api-key": GEMINI_API_KEY
        }
        data = {
            "contents": [
                {
                    "parts": [
                        {
                            "text": "你是一名幼儿园群聊管理助手，核心功能为对接收到的群聊消息进行内容分析与分类，具体规则如下："
                                    "分类结果：仅分为 \"重要消息\" 和 \"普通消息\" 两类。"
                                    "重要消息判定标准：满足以下任一条件即归为重要消息"
                                    "消息内容包含 @符号（如 @老师、@全体成员）"
                                    "家长吐槽类内容（如表达对幼儿园餐食、活动安排的不满）"
                                    "家长夸赞类内容（如表扬老师照顾、肯定幼儿园活动）"
                                    "家长提问类内容（无需带问号，根据句型判断，如 \"请问明天需要带什么材料\"\"下周家长会时间是哪天\"）"
                                    "家长投诉类内容（如明确表达对幼儿园服务、管理的正式不满与诉求）"
                                    "普通消息判定标准：除上述重要消息外，其余均为普通消息，包括家长间闲聊、日常交流（如 \"今天天气真好\"\"孩子今天很开心\"）。"
                                    "请基于此规则，对输入的每一条群聊消息（系统消息（@system）除外），输出 \"分类结果\"\n\n"
                                    f"消息内容: {content}\n请回答'重要'或'普通'。"
                        }
                    ]
                }
            ],
            "generationConfig": {
                "temperature": 0.2,
                "maxOutputTokens": 100
            }
        }
        resp = requests.post(GEMINI_API_URL, headers=headers, json=data, timeout=15)
        resp.raise_for_status()
        result = resp.json()
        reply = result["candidates"][0]["content"]["parts"][0]["text"].strip()

        return "重要" in reply
    except Exception as e:
        print(f"[AI ERROR] 分析消息失败: {e}")
        raise


@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def analyze_message_with_retry(content: str) -> bool:
    """
    带重试机制的AI分析函数
    """
    time.sleep(random.uniform(1.0, 3.0))
    return analyze_message(content)


def batch_analyze_messages():
    """批量分析过去一小时内的消息"""
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()

    try:
        # 分析过去一小时内的消息
        start_time = (datetime.now() - timedelta(hours=1)).strftime("%Y-%m-%d %H:%M:%S")
        end_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        print(f"[INFO] 分析最近一小时的消息: {start_time} 到 {end_time}")

        # 获取消息
        cur.execute("""
            SELECT id, chat_name, sender, content, msg_time, raw 
            FROM messages 
            WHERE msg_time BETWEEN ? AND ?
        """, (start_time, end_time))

        messages_to_analyze = cur.fetchall()

        if not messages_to_analyze:
            print(f"[INFO] 该时间段内没有消息")
            return

        print(f"[INFO] 开始分析 {len(messages_to_analyze)} 条消息")

        # 逐条分析消息
        important_count = 0
        for i, msg in enumerate(messages_to_analyze):
            msg_id, chat_name, sender, content, msg_time, raw_json = msg

            # 每处理5条消息休息一下
            if i > 0 and i % 5 == 0:
                rest_time = random.uniform(2, 5)
                print(f"[INFO] 已处理 {i} 条消息，休息 {rest_time:.1f} 秒...")
                time.sleep(rest_time)

            try:
                important = analyze_message_with_retry(content)

                if important:
                    raw_data = json.loads(raw_json)
                    save_message(chat_name, sender, content, msg_time, raw_data, important=True)
                    important_count += 1
                    time.sleep(random.uniform(0.5, 1.5))

            except Exception as e:
                print(f"[ERROR] 分析消息 {msg_id} 失败，跳过: {e}")
                continue

        print(f"[INFO] 处理完成，发现 {important_count} 条重要消息")

    except Exception as e:
        print(f"[ERROR] 批量分析失败: {e}")
    finally:
        conn.close()


# ============ 工具 ============
def now_str() -> str:
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def safe_get(obj, names: list, default=""):
    for n in names:
        try:
            v = getattr(obj, n)
            if v is not None:
                return v
        except Exception:
            continue
    return default


def get_group_name(msg):
    """获取消息的群聊名称，兼容不同的消息对象结构"""
    for attr in ["member", "group", "Group", "chat", "Chat"]:
        try:
            value = getattr(msg, attr, None)
            if value:
                return value
        except Exception:
            pass
    return "未知群聊"


def batch_processor():
    """后台批量处理线程 - 严格每小时处理一次"""
    print("[INFO] 批量处理线程启动，每小时处理一次")

    while not stop_event.is_set():
        try:
            # 严格等待1小时
            print(f"[INFO] 等待下一次处理（1小时后）...")
            stop_event.wait(3600)

            if stop_event.is_set():
                break

            print(f"[INFO] 开始每小时批量处理 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            batch_analyze_messages()

        except Exception as e:
            print(f"[ERROR] 批量处理线程异常: {e}")
            time.sleep(60)


# ============ 主程序 ============
def main():
    init_db(DB_PATH)
    wx = WeChat()

    # 启动批量处理线程
    batch_thread = threading.Thread(target=batch_processor, daemon=True)
    batch_thread.start()

    for chat_name in TARGET_CHATS:
        try:
            wx.ChatWith(chat_name)
        except Exception:
            pass

    print(f"[INFO] 对话监听绑定：{', '.join(TARGET_CHATS)}")

    for chat_name in TARGET_CHATS:
        def on_message(msg, chat, chat_name=chat_name):
            human_like_delay()

            try:
                # 获取消息ID
                msg_id = getattr(msg, "id", None)
                if not msg_id:
                    msg_content = safe_get(msg, ["content", "Content", "text", "Text", "raw", "msg"], default="")
                    msg_id = hash(f"{chat_name}_{msg_content}")

                # 检查重复消息
                if chat_name in last_messages and last_messages[chat_name] == msg_id:
                    return

                last_messages[chat_name] = msg_id

                ts = now_str()
                sender = getattr(msg, "sender", "") or getattr(msg, "Sender", "") or "未知用户"
                content = safe_get(msg, ["content", "Content", "text", "Text", "raw", "msg"], default="")
                actual_chat_name = chat_name or get_group_name(msg)

                raw = {}
                try:
                    raw = msg.__dict__ if hasattr(msg, "__dict__") else {}
                except Exception:
                    raw = {"repr": repr(msg)}

                print(f"[{ts}]<{actual_chat_name}>@{sender}：{content}")

                # 只保存到普通消息表，不立即分析
                save_message(actual_chat_name, sender, content, ts, raw)

            except Exception as e:
                print(f"[ERROR] 回调异常: {e}")

        wx.AddListenChat(chat_name, on_message)

    print(f"[INFO] 已开始监听群聊：{', '.join(TARGET_CHATS)}")
    print("[INFO] 按 Ctrl+C 退出程序")

    def handle_sigint(sig, frame):
        print("\n[INFO] 正在停止程序...")
        stop_event.set()

    signal.signal(signal.SIGINT, handle_sigint)

    # 主循环
    while not stop_event.is_set():
        stop_event.wait(0.5)

    # 程序结束前进行最后一次处理
    print("[INFO] 程序结束，进行最后一次批量处理...")
    batch_analyze_messages()
    print("[INFO] 已停止。")


if __name__ == "__main__":
    main()